{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "0829c265-2443-4b50-b26a-6f53d68e8933",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import glob\n",
    "import shutil\n",
    "import numpy as np\n",
    "import os.path as osp\n",
    "\n",
    "from tqdm import tqdm\n",
    "from ultralytics import YOLO\n",
    "from functools import partial\n",
    "from pycocotools.coco import COCO\n",
    "from multiprocessing import Pool, cpu_count\n",
    "\n",
    "from my_class import *\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "21fc9dd0-d7cd-41f0-9e59-a1808b682389",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/data/pt/data/NikeDatasets/ReviewData/update_data/0424/no_labels_xiantou/xiantou/images\n"
     ]
    }
   ],
   "source": [
    "flaw_name = \"xiantou\"\n",
    "defect_name = [\"xiantou\"]\n",
    "\n",
    "date = \"0424\"\n",
    "daily_dir = f\"/data/pt/data/NikeDatasets/ReviewData/{date}/no_labels_xiantou\"\n",
    "\n",
    "result_dir = f\"/data/pt/data/NikeDatasets/ReviewData/update_data/{date}/no_labels_xiantou/{flaw_name}\"\n",
    "\n",
    "result_dir_images = osp.join(result_dir, \"images\")\n",
    "result_dir_labels = osp.join(result_dir, \"labels\")\n",
    "\n",
    "if not osp.exists(result_dir_images):\n",
    "    os.makedirs(result_dir_images)\n",
    "    os.makedirs(result_dir_labels)\n",
    "    os.makedirs(osp.join(result_dir, \"annotations\"))\n",
    "    \n",
    "footFind = YOLO(\"/home/user/workspace/model_zoo/nike_model/FootFind.onnx\", task=\"detect\")\n",
    "print(result_dir_images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "51e7a33e-0418-4b0e-a330-ca4dae713d74",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_path(root, img):\n",
    "    img_path = glob.glob(f\"{root}/**/{img['file_name']}\", recursive=True)\n",
    "    if len(img_path) != 1:\n",
    "        print(f\"⚠️ Image not found: {root} {img['file_name']}\")\n",
    "        return None\n",
    "    if img['height'] == 0 or img['width'] == 0:\n",
    "        print(f\"⚠️ Image Shape Incorrect: {root} {img['file_name']}\")\n",
    "        return None\n",
    "    return img_path[0]\n",
    "\n",
    "def get_ann_dir(root):\n",
    "    i = 3\n",
    "    while i > 0:\n",
    "        res = osp.join(root, f\"annotations_v{i}\")\n",
    "        if osp.exists(res):\n",
    "            break\n",
    "        i -= 1\n",
    "    if not osp.exists(res):\n",
    "        res = osp.join(root, \"annotations\")\n",
    "    \n",
    "    return res\n",
    "\n",
    "def expand_coordinate(bbox, offset, ori_shape):\n",
    "    x, y, w, h = list(map(int, bbox))\n",
    "    ori_width= ori_shape[1]\n",
    "    ori_height =ori_shape[0]\n",
    "\n",
    "    x1 = max(0, x-offset*ori_width)\n",
    "    x2 = min(x+w+offset*ori_width, ori_width)\n",
    "    y1 = max(0, y-offset*ori_height)\n",
    "    y2 = min(y+h+offset*ori_height, ori_height)\n",
    "    return list(map(int, (x1, y1, x2-x1, y2-y1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "dcc7bb53-9c70-4363-833a-7340eebe2fdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_one_image(img, coco, img_root_dir, name2id, defect_name, result_dir_labels, result_dir_images, is_expand, save_label, save_img, save_empty):\n",
    "    try:\n",
    "        img_path = get_path(img_root_dir, img)\n",
    "        img_w, img_h = img[\"width\"], img[\"height\"]\n",
    "        if save_label:\n",
    "            # Step1: 获取ROI区域\n",
    "            if \"drunk\" in name2id:\n",
    "                base = coco.loadAnns(coco.getAnnIds(imgIds=img['id'], catIds=[name2id[\"drunk\"]]))\n",
    "            else:\n",
    "                base = []\n",
    "    \n",
    "            if len(base) == 0:\n",
    "                res = footFind(img_path, imgsz=320, device=\"cpu\",verbose=False)\n",
    "                x, y, x2, y2 = res[0].boxes.xyxy.numpy()[0]\n",
    "                roi = [x, y, x2 - x, y2 - y]\n",
    "            else:\n",
    "                roi = base[0]['bbox']\n",
    "    \n",
    "            if is_expand:\n",
    "                roi = expand_coordinate(roi, 0.02, (img_h, img_w))\n",
    "            roi = BBox.from_xywh(*roi)\n",
    "    \n",
    "            # Step2: 保存 label\n",
    "\n",
    "            target_ids = [name2id[cat] for cat in defect_name if cat in name2id]\n",
    "            annotations = coco.loadAnns(coco.getAnnIds(imgIds=img[\"id\"], catIds=target_ids))\n",
    "\n",
    "            txt = []\n",
    "            for ann in annotations:\n",
    "                category_id = int(ann['category_id'])\n",
    "                coordinates = np.array(ann['segmentation']).flatten().tolist()\n",
    "                mask = PolyMask(coordinates).normalize_in_bbox(roi)\n",
    "                txt.append(\"0 \" + \" \".join(map(str, mask.flat())))\n",
    "\n",
    "            dist = result_dir_labels\n",
    "            os.makedirs(dist, exist_ok=True)\n",
    "            with open(os.path.join(dist, os.path.basename(img_path).replace('.png', '.txt')), \"w\") as f:\n",
    "                for annotation in txt:\n",
    "                    f.write(annotation + \"\\n\")\n",
    "\n",
    "        # Step3: 保存图像\n",
    "        if save_img:\n",
    "            if not save_empty and save_label and len(txt) == 0:\n",
    "                return\n",
    "            x1, y1, x2, y2 = list(map(int, roi.xyxy))\n",
    "            img_array = cv2.imread(img_path)[y1:y2, x1:x2]\n",
    "            os.makedirs(result_dir_images, exist_ok=True)\n",
    "            cv2.imwrite(os.path.join(result_dir_images, os.path.basename(img_path)), img_array)\n",
    "    except Exception as e:\n",
    "        print(f\"[Error] Processing image id={img['id']} failed: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "eff15a8e-8b09-458e-ab5b-e2278bf93543",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/user/workspace/datasets/NikeDatasets/ReviewData/0424/no_labels_xiantou/annotations\n"
     ]
    }
   ],
   "source": [
    "# rootDir = '/home/user/workspace/datasets/NikeDatasets/ReviewData/0424/filtered_xiantou'\n",
    "rootDir = '/home/user/workspace/datasets/NikeDatasets/ReviewData/0424/no_labels_xiantou'\n",
    "img_root_dir = osp.join(rootDir, \"images\")\n",
    "anno_dir = get_ann_dir(rootDir)\n",
    "print(anno_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "dcd65072-5755-461b-8dd2-af9c253bb1db",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading annotations into memory...\n",
      "Done (t=0.00s)\n",
      "creating index...\n",
      "index created!\n",
      "{'线头': 1}\n"
     ]
    }
   ],
   "source": [
    "save_img = True\n",
    "save_lable = True\n",
    "crop = True\n",
    "is_expand = True\n",
    "\n",
    "save_empty = True\n",
    "\n",
    "json_file = glob.glob(f\"{anno_dir}/**/*.json\", recursive=True)\n",
    "coco = COCO(json_file[0])\n",
    "categories = coco.loadCats(coco.getCatIds())\n",
    "\n",
    "name2id = {cat[\"name\"]:cat[\"id\"] for cat in categories}\n",
    "print(name2id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "32833048-4e24-44ae-b9a5-1aa544d8e140",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/36 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading /home/user/workspace/model_zoo/nike_model/FootFind.onnx for ONNX Runtime inference...\n",
      "Loading /home/user/workspace/model_zoo/nike_model/FootFind.onnx for ONNX Runtime inference...\n",
      "Loading /home/user/workspace/model_zoo/nike_model/FootFind.onnx for ONNX Runtime inference...\n",
      "Loading /home/user/workspace/model_zoo/nike_model/FootFind.onnx for ONNX Runtime inference...\n",
      "Loading /home/user/workspace/model_zoo/nike_model/FootFind.onnx for ONNX Runtime inference...\n",
      "Loading /home/user/workspace/model_zoo/nike_model/FootFind.onnx for ONNX Runtime inference...\n",
      "Loading /home/user/workspace/model_zoo/nike_model/FootFind.onnx for ONNX Runtime inference...\n",
      "Using ONNX Runtime CPUExecutionProvider\n",
      "Using ONNX Runtime CPUExecutionProvider\n",
      "Using ONNX Runtime CPUExecutionProvider\n",
      "Using ONNX Runtime CPUExecutionProvider\n",
      "Using ONNX Runtime CPUExecutionProvider\n",
      "Using ONNX Runtime CPUExecutionProvider\n",
      "Using ONNX Runtime CPUExecutionProvider\n",
      "Loading /home/user/workspace/model_zoo/nike_model/FootFind.onnx for ONNX Runtime inference...\n",
      "Using ONNX Runtime CPUExecutionProvider\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 36/36 [00:12<00:00,  2.84it/s]\n"
     ]
    }
   ],
   "source": [
    "img_list = list(coco.imgs.values())\n",
    "num_workers = min(8, cpu_count())  # 你可以根据自己机器的核数调整\n",
    "\n",
    "with Pool(num_workers) as pool:\n",
    "    with tqdm(total=len(img_list)) as pbar:\n",
    "        for _ in pool.imap_unordered(partial(\n",
    "            process_one_image,\n",
    "            coco=coco,\n",
    "            img_root_dir=img_root_dir,\n",
    "            name2id=name2id,\n",
    "            defect_name=defect_name,\n",
    "            result_dir_labels=result_dir_labels,\n",
    "            result_dir_images=result_dir_images,\n",
    "            is_expand=is_expand,\n",
    "            save_label=save_lable,\n",
    "            save_img=save_img,\n",
    "            save_empty=save_empty\n",
    "        ), img_list):\n",
    "            pbar.update(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1242a92c-bd77-4cff-b92e-41e96b65917c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50db0f72-8885-4352-8024-92bfafd0c548",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
